import random
from functools import partial

import numpy as np
import torch

class PatchClipper:
    
    def __init__(self, **kwargs):

        self.crop_strategy = {
            "shift template" : self._crop_shifted_template_patchs, 
            "center same" : self._crop_same_patchs_center, 
            "shift template location" : self._crop_shifted_template_location, 
            "None" : lambda a, b : (a, b), 
        }
        
        self.clipper = kwargs["name"]
        kwargs.pop("name")

        if self.clipper in (
            "shift template", "shift template location" 
        ):
            self.search_domain = "A"
            if "search_domain" in kwargs:
                self.search_domain = kwargs["search_domain"]
            if self.search_domain == "":
                self.search_domain = "A"
                kwargs["search_domain"] = "A"

        self.crop = partial(
            self.crop_strategy[self.clipper], **kwargs
        )

    def _crop_shifted_template_patchs(
        self, patch_a, patch_b, *, template_size=144, 
        shift_x=0, shift_y=0, search_rad=32, search_domain="A", 
        ret_shift=True, fixed=False, mask=None
    ):
        """
        crop patchs and call get_heatmap acrroding to shift x, y

        Param
        -----
        ret_shift   :   return tensor of shift x, y
        fixed       :   move template with (shift_x, ..y) when True,
                        or the random of (-shift_x~shift_x, ..y)
        """

        crop_size = template_size + 2 * search_rad

        search_mask, template_mask = None, None
    
        if search_domain == 'A':
            search_img = patch_a
            template_img = patch_b
            search_mask = mask
        else:
            search_img = patch_b
            template_img = patch_a
            template_mask = mask
        ms = search_rad

        c, h, w = search_img.size()

        crop_begin_h, crop_begin_w = h//2 - crop_size//2, w//2 - crop_size//2
        crop_end_h, crop_end_w = crop_begin_h+crop_size, crop_begin_w+crop_size

        if not fixed:
            s_x = random.randint(-shift_x, shift_x)
            s_y = random.randint(-shift_y, shift_y)
        else:
            s_x = shift_x
            s_y = shift_y

        search_patch = search_img[:, 
            crop_begin_h:crop_end_h, 
            crop_begin_w:crop_end_w
        ]
        template_patch = template_img[:, 
            crop_begin_h+s_y+ms : crop_end_h+s_y-ms, 
            crop_begin_w+s_x+ms : crop_end_w+s_x-ms
        ]
        if not search_mask is None:
            search_mask = search_mask[:, 
            crop_begin_h:crop_end_h, 
            crop_begin_w:crop_end_w
        ]
        if not template_mask is None:
            template_mask = template_mask[:, 
            crop_begin_h+s_y+ms : crop_end_h+s_y-ms, 
            crop_begin_w+s_x+ms : crop_end_w+s_x-ms
        ]

        gt_map = self._get_heatmap(s_x, s_y, search_rad=search_rad)
        # gt_map = self._get_heatmap_gaussian(s_x, s_y, search_rad=search_rad)

        if search_domain == 'A':
            if not search_mask is None:
                patch_a = (search_patch, search_mask)
            else:
                patch_a = search_patch
            patch_b = template_patch
        else:
            if not template_mask is None:
                patch_a = (template_patch, template_mask)
            else:
                patch_a = template_patch
            patch_b = search_patch

        if ret_shift:
            return patch_a, patch_b, gt_map, torch.tensor([search_rad+s_x, search_rad+s_y])
        else:
            return patch_a, patch_b, gt_map

    def _get_heatmap(self, shift_x, shift_y, *, 
        w=0, h=0, search_rad=32, extent_pos=1
    ):
        """
        get ground truth score map of match
        """
        if w==0 and h==0 and search_rad==0:
            print("wrong: all of w, h, search_rad are 0 !")
            return None

        if w==0 or h==0:
            w = 2 * search_rad + 1
            h = 2 * search_rad + 1
        
        if search_rad == 0:
            search_rad = w // 2

        gt_map = torch.zeros(h, w)

        x = shift_x
        y = shift_y
        
        gt_map[
            search_rad-extent_pos+y : search_rad+extent_pos+y+1, 
            search_rad-extent_pos+x : search_rad+extent_pos+x+1
        ] = 1

        return gt_map.float()

    def _get_heatmap_gaussian(self, shift_x=0, shift_y=0, *, 
        search_rad=32, sigma=1.0, keep_val=3
    ):

        size = 2 * search_rad + 1
        K = np.zeros((size, size), dtype=np.float)

        pad = search_rad

        for x in range(-pad, -pad + size):
            for y in range(-pad, -pad + size):
                if not (
                    ((y+pad+shift_y) >= 0 and (y+pad+shift_y) < size)
                and ((x+pad+shift_x) >= 0 and (x+pad+shift_x) < size)
                ):
                    continue
                K[y+pad+shift_y, x+pad+shift_x] \
                    = np.exp( -(x ** 2 + y ** 2) / (2 * (sigma ** 2)))

        if keep_val != -1:
            rad = keep_val * sigma
            thresh = np.exp( -(rad ** 2 + rad ** 2) / (2 * (sigma ** 2)))
            K[K<thresh] = 0

        # K /= (2 * np.pi * sigma * sigma)
        # K /= K.sum()

        return torch.tensor(K)

    def _crop_same_patchs_center(
        self, patch_a, patch_b, *, template_size=None
    ):
        """
        crop patchs at center with the same template size

        Param
        -----
        template_size   :   size of template
        """

        c, h, w = patch_a.size()

        ts = template_size
        crop_begin_h, crop_begin_w = h//2 - ts//2, w//2 - ts//2
        crop_end_h, crop_end_w = crop_begin_h+ts, crop_begin_w+ts

        patch_a = patch_a[:, 
            crop_begin_h:crop_end_h,
            crop_begin_w:crop_end_w,
        ]

        patch_b = patch_b[:, 
            crop_begin_h:crop_end_h,
            crop_begin_w:crop_end_w,
        ]

        return patch_a, patch_b

    def _crop_shifted_template_location(
        self, patch_a, patch_b, *, template_size=144, 
        shift_x=0, shift_y=0, search_rad=32, search_domain="A", 
        ret_shift=True, fixed=False
    ):
        crop_size = template_size + 2 * search_rad
    
        if search_domain == 'A':
            search_img = patch_a
            template_img = patch_b
        else:
            search_img = patch_b
            template_img = patch_a
        ms = search_rad

        c, h, w = search_img.size()

        crop_begin_h, crop_begin_w = h//2 - crop_size//2, w//2 - crop_size//2
        crop_end_h, crop_end_w = crop_begin_h+crop_size, crop_begin_w+crop_size

        if not fixed:
            s_x = random.randint(-shift_x, shift_x)
            s_y = random.randint(-shift_y, shift_y)
        else:
            s_x = shift_x
            s_y = shift_y

        # x_left_top, y_left_top, x_right_down, y_right_down, 
        search_loc = torch.tensor([
            crop_begin_w, crop_begin_h, 
            crop_end_w, crop_end_h, 
        ]).float()
        template_loc = torch.tensor([
            crop_begin_w+s_x+ms, crop_begin_h+s_y+ms, 
            crop_end_w+s_x-ms, crop_end_h+s_y-ms, 
        ]).float()

        gt_map = self._get_heatmap(s_x, s_y, search_rad=search_rad)

        if search_domain == 'A':
            loc_a = search_loc
            loc_b = template_loc
        else:
            loc_a = template_loc
            loc_b = search_loc

        if ret_shift:
            return loc_a, loc_b, gt_map, torch.tensor([search_rad+s_x, search_rad+s_y])
        else:
            return loc_a, loc_b, gt_map


if __name__ == "__main__":
    from dotmap import DotMap

    # mode_ = "shift template"
    mode_ = "shift template location"

    clipper = DotMap({
        "name" : mode_,
        "template_size" : 144,
        "shift_x" : 10,
        "shift_y" : 10,
        "search_rad" : 32,
        "search_domain" : "B", 
        "fixed": True
    })

    pc = PatchClipper(**(clipper.toDict()))

    a = torch.ones((1, 256, 256)).float().cuda()
    b = torch.zeros((1, 256, 256)).float().cuda()

    sar_p, opt_p, gt_map, shift = pc.crop(a, b)

    if mode_ == "shift template":
        print(sar_p.shape, opt_p.shape)
        print(torch.argmax(gt_map, dim=1), torch.argmax(gt_map, dim=0), gt_map.shape, shift)
    else:
        print(sar_p, opt_p)
        print(torch.argmax(gt_map, dim=1), torch.argmax(gt_map, dim=0), gt_map.shape, shift)

